#!/usr/bin/env python
# coding=utf-8
# @FileName :main.py
# @Time     :2023/3/27 19:25
# @Author   :ZXC
# @Desc     :入口文件
import time
from datetime import datetime
from json import JSONDecodeError
from pathlib import Path
from typing import Union
from urllib.parse import unquote

import urllib3
from loguru import logger
from magic import magic
from requests import Request, Response, Session
from requests.exceptions import MissingSchema, InvalidSchema, InvalidURL, RequestException

from .schemas import ResponseBaseModel, HttpClientResponseModel, ResponseStatModel

# 清除InsecureRequestWarning告警
urllib3.disable_warnings()


class HttpResponse(Response):
    """
    HTTP请求响应对象
    """

    def raise_for_status(self):
        """
        抛出响应的错误信息error字段
        :return:
        """
        if hasattr(self, "error") and self.error:
            raise self.error
        Response.raise_for_status(self)


class HttpSession(Session):
    """
    执行HTTP请求，并在请求之间保存session、cookies。
    """

    def __init__(self):
        """
        初始化
        """
        logger.debug('初始化HttpSession对象')
        super(HttpSession, self).__init__()
        # 初始化响应数据
        self.__init_response_meta()

    def __init_response_meta(self):
        """
        初始化响应模型
        @return:
        """
        self.response_model = HttpClientResponseModel(
            response=ResponseBaseModel(),
            error='',
            stat=ResponseStatModel()
        )

    def request(self, method: str, url: str, **kwargs) -> HttpClientResponseModel:
        """
        发送请求
        @param method: 请求方法
        @param url: 请求url
        @return:
        """
        logger.debug('准备初始化响应模型并执行接口请求')
        # 初始化响应数据
        self.__init_response_meta()

        # 设置请求默认的超时时间(s)
        kwargs.setdefault("timeout", 120)

        # 记录执行开始时间
        start_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.response_model.stat.start_time = start_time

        # 发送请求
        response: Union[Response, HttpResponse] = self._send_request_safe_mode(method, url, **kwargs)

        # 记录请求和响应，不包括30X重定向
        self._get_response_records(response)

        logger.debug('检查接口请求是否异常')
        # 抛出响应异常
        try:
            response.raise_for_status()
        except RequestException as e:
            # 记录异常
            self.response_model.error = str(e)
            logger.error("检查接口请求异常：{}".format(str(e)))
        # 记录执行结束时间
        end_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.response_model.stat.end_time = end_time
        logger.info(f'接口请求完成，最终返回结果：\n{self.response_model.dict(exclude={"text", "content"})}')
        return self.response_model

    def _send_request_safe_mode(self, method: str, url: str, **kwargs) -> Union[Response, HttpResponse]:
        """
        发送一个HTTP请求，并捕获由于连接问题可能发生的任何异常。
        @param method: 请求方法
        @param url: 请求URL
        @param kwargs: 请求的其他关键词参数
        @return: Response对象或其子类对象HttpResponse
        """
        logger.debug('执行接口请求')
        try:
            self.stream = kwargs.get('stream', False)
            return Session.request(self, method=method, url=url, **kwargs)
        except (MissingSchema, InvalidSchema, InvalidURL):
            raise
        except RequestException as e:
            r = HttpResponse()
            r.error = e  # 封装异常到响应对象中
            r.status_code = 0  # 响应状态码为0时，content为None
            r.request = Request(method, url).prepare()  # 封装请求对象到响应对象中
            return r

    def _get_response_records(self, resp_obj: Union[Response, HttpResponse]):
        """
        从Response()对象中获取响应的信息
        @param resp_obj: 处理后的响应dict
        @return:
        """
        logger.info('接口请求结束，准备处理接口返回数据')
        try:
            self.response_model.response.ok = resp_obj.ok
        except Exception as e:
            self.response_model.response.ok = False
            logger.warning('接口请求失败，数据未处理。')
            logger.debug(e)
        else:
            for k, v in resp_obj.headers.items():
                if k.lower() == 'content-type':
                    self.response_model.response.content_type = v
                    break
            self.response_model.response.status_code = resp_obj.status_code
            self.response_model.response.reason = resp_obj.reason
            self.response_model.response.elapsed = resp_obj.elapsed.total_seconds()
            self.response_model.response.url = resp_obj.url
            self.response_model.response.encoding = resp_obj.encoding
            self.response_model.response.apparent_encoding = str(resp_obj.apparent_encoding)
            self.response_model.response.is_redirect = resp_obj.is_redirect
            self.response_model.response.is_permanent_redirect = resp_obj.is_permanent_redirect
            self.response_model.response.cookies = dict(resp_obj.cookies)
            self.response_model.response.headers = dict(resp_obj.headers)
            try:
                # 记录响应的json数据
                self.response_model.response.json_data = resp_obj.json()
            except JSONDecodeError:
                if "image" in self.response_model.response.content_type:
                    self.response_model.response.content = resp_obj.content
                else:
                    self.response_model.response.text = resp_obj.text
            self.response_model.response.iter_lines = resp_obj.iter_lines
            self.response_model.response.iter_content = resp_obj.iter_content
            self.response_model.response.links = resp_obj.links
            self.response_model.response.history = resp_obj.history

            logger.info('接口请求成功，返回数据处理完成')

    @staticmethod
    def _get_download_filename(response: HttpClientResponseModel) -> str:
        """
        获取需要下载的文件的文件名
        @param response:
        @return:
        """

        def _get_file_name_and_type_by_content_disposition(inner_response: HttpClientResponseModel) -> tuple:
            """
            从Content-Disposition提取值为attachment的值(文件名和文件类型)
            @param inner_response:
            @return:
            """
            logger.debug('从Content-Disposition提取值为attachment的值(文件名和文件类型)')
            for k, v in inner_response.response.headers.items():
                if k.lower() == 'content-disposition':
                    content_disposition: str = v
                    if content_disposition.startswith('attachment;'):
                        filename = content_disposition.split(';')[1]
                        full_name = filename.rsplit('=')[-1].strip('"')
                        # 解决中文乱码
                        name = full_name.rsplit('.', maxsplit=1)[0].encode('ISO-8859-1').decode('utf8')
                        file_type = full_name.rsplit('.')[-1]
                        # 解决url编码问题
                        try:
                            name = unquote(name, 'utf-8')
                        except Exception as e:
                            logger.debug(e)
                            pass
                        logger.debug(f'提取结果【{name}.{file_type}】')
                        return name, file_type
            logger.debug(f'从Content-Disposition未提取到')
            return None, None

        def _get_file_type_by_content_type(inner_response: HttpClientResponseModel) -> Union[str, None]:
            """
            从content_type提取文件类型
            @param inner_response:
            @return:
            """
            logger.debug('从content_type提取文件类型')
            file_type = None
            inner_content_type = inner_response.response.content_type.split(';', maxsplit=1)[0]
            if inner_content_type.startswith('text/'):
                if 'plain' in inner_content_type:
                    file_type = None
                elif 'xml' in inner_content_type:
                    file_type = 'xml'
                elif 'html' in inner_content_type:
                    file_type = 'html'
                elif 'javascript' in inner_content_type or 'ecmascript' in inner_content_type or \
                        'jscript' in inner_content_type or 'livescript' in inner_content_type:
                    file_type = 'js'
                elif 'css' in inner_content_type:
                    file_type = 'css'
            elif inner_content_type.startswith('image/'):
                if 'x-icon' in inner_content_type or 'vnd.microsoft.icon' in inner_content_type:
                    file_type = 'ico'
            elif inner_content_type.startswith('audio/'):
                pass
            elif inner_content_type.startswith('video/'):
                pass
            elif inner_content_type.startswith('application/'):
                if 'json' in inner_content_type:
                    file_type = 'json'
                elif 'javascript' in inner_content_type or 'ecmascript' in inner_content_type:
                    file_type = 'js'
                elif 'pdf' in inner_content_type:
                    file_type = 'pdf'
                elif 'xml' in inner_content_type:
                    file_type = 'xml'
            logger.debug(f'从content_type提取到文件类型：{file_type}')
            return file_type

        def _get_file_type_by_file_stream(inner_response: HttpClientResponseModel, chunk_size: int = 2048):
            """
            文件流分析得到文件类型
            @param inner_response: 响应对象
            @param chunk_size: 单次读取最大行数
            @return:
            """
            logger.debug('从文件流分析文件类型')
            try:
                for data in inner_response.response.iter_content(chunk_size=chunk_size):  # 从流中读取，建议不少于2048字节
                    res = magic.from_buffer(data, mime=True).split('/')[-1]
                    if res == 'plain':
                        res = 'txt'
                    logger.debug(f'从文件流分析文件类型：{res}')
                    return res
            except:
                return None

        file_name, file_type = _get_file_name_and_type_by_content_disposition(response)
        # 从content_disposition没解析到文件名和文件类型，尝试从content_type解析
        if file_name is None and file_type is None:
            file_type = _get_file_type_by_content_type(response)
        # 从content_disposition和content_type没解析到文件类型，尝试从文件流解析
        if file_type is None:
            file_type = _get_file_type_by_file_stream(response)
        if file_name is None:
            file_name = f'response{int(round(time.time()) * 1000)}'
        # 最终的文件名
        final_file_name = f'{file_name}'
        if file_type:
            final_file_name += f'.{file_type}'

        logger.debug(f'解析到: 文件【{final_file_name}】')
        return final_file_name

    @staticmethod
    def down_load(response: HttpClientResponseModel, save_dir: str | Path = '.', filename: str = None,
                  chunk_size: int = 64 * 1024):
        """
        文件下载器
        @param response: 响应对象
        @param save_dir: 存储路径
        @param filename: 存储名称
        @param chunk_size: 单次读取最大行数
        @return:
        """

        # 判断是否自定义名称
        if filename is None:
            filename = HttpSession._get_download_filename(response)
        final_path = Path(save_dir).joinpath(filename)
        logger.debug(f'准备保存文件【{final_path}】')
        with open(final_path, "wb") as file:
            for data in response.response.iter_content(chunk_size=chunk_size):
                if not data:
                    break
                file.write(data)

        logger.info(f'保存文件完成【{final_path}】')
